#!/usr/bin/env python3
import rospy
import sensor_msgs.point_cloud2 as pc2
from sensor_msgs.msg import PointCloud2, Image, CompressedImage
from nav_msgs.msg import Odometry
from geometry_msgs.msg import PointStamped
import numpy as np
import math
import cv2
import cv2.aruco as aruco
from cv_bridge import CvBridge
from std_msgs.msg import String, Int32
from parse_goal_direction import parse_goal_direction
from parse_goal_direction import validate_direction

class FrontierExplorer:
    def __init__(self):
        rospy.init_node('frontier_explorer', anonymous=True)
        self.frontiers = []
        self.robot_pos = None
        self.current_goal = None
        
        # 从ROS参数服务器读取参数
        self.goal_reached_threshold = rospy.get_param('~goal_reached_threshold', 0.5)
        
        # 探索状态
        self.exploration_active = False
        self.consecutive_empty_frontiers = 0  # 连续没有frontier的次数
        self.max_empty_count = rospy.get_param('~max_empty_count', 5)
        
        # 目标点发布频率设置
        self.waypoint_publish_rate = 10.0  # 发布频率，单位：Hz
        self.last_waypoint_time = rospy.Time.now()
        
        # 探索区域边界（相对于机器人起始位置）
        # 从ROS参数服务器读取边界参数
        self.exploration_bounds = {
            'x_min': rospy.get_param('~exploration_x_min', -50.0),
            'x_max': rospy.get_param('~exploration_x_max', 50.0),
            'y_min': rospy.get_param('~exploration_y_min', -50.0),
            'y_max': rospy.get_param('~exploration_y_max', 50.0),
        }
        
        # 打印从参数服务器读取的参数信息
        rospy.loginfo("=" * 50)
        rospy.loginfo("Frontier Explorer 参数配置:")
        rospy.loginfo(f"  目标到达阈值: {self.goal_reached_threshold}")
        rospy.loginfo(f"  最大空frontier计数: {self.max_empty_count}")
        rospy.loginfo(f"  航点发布频率: {self.waypoint_publish_rate} Hz")
        rospy.loginfo(f"  探索区域边界:")
        rospy.loginfo(f"    X: [{self.exploration_bounds['x_min']:.1f}, {self.exploration_bounds['x_max']:.1f}]")
        rospy.loginfo(f"    Y: [{self.exploration_bounds['y_min']:.1f}, {self.exploration_bounds['y_max']:.1f}]")
        rospy.loginfo("=" * 50)
        
        # 机器人起始位置（用于计算相对坐标）
        self.robot_start_pos = None
        self.robot_yaw = 0.0  # 机器人朝向角度
        
        # 方向偏好设置，可以通过参数或服务动态修改
        self.preferred_direction = "left"  # 可选: "right front", "left front", "front", "back", "right back", "left back"
        
        # 方向角度定义（相对于机器人前方）
        self.direction_angles = {
            "front": 0,
            "left front": 45,
            "left": 90,
            "left back": 135,
            "back": 180,
            "right back": 225,
            "right": 270,
            "right front": 315
        }
        
        # 方向权重（越接近偏好方向权重越高）
        self.direction_weight = 2.0

        # Aruco检测相关
        self.bridge = CvBridge()
        self.aruco_dict = aruco.Dictionary_get(aruco.DICT_4X4_50)
        self.aruco_params = aruco.DetectorParameters_create()
        self.aruco_detected = False
        self.aruco_position = None
        
        # 摄像头内参 (根据URDF文件计算)
        image_width = 1080  # 从URDF文件更新
        image_height = 720  # 从URDF文件更新
        horizontal_fov = 2.0  # 弧度，来自URDF
        
        # 计算焦距
        fx = (image_width / 2.0) / math.tan(horizontal_fov / 2.0)
        fy = fx  # 假设像素是正方形
        cx = image_width / 2.0  # 主点x坐标
        cy = image_height / 2.0  # 主点y坐标
        
        self.camera_matrix = np.array([[fx, 0.0, cx],
                                      [0.0, fy, cy],
                                      [0.0, 0.0, 1.0]])
        self.dist_coeffs = np.array([0.0, 0.0, 0.0, 0.0, 0.0])
        self.aruco_size = 1  # Aruco标记的实际大小(米)
        
        # 相机外参 (相机相对于机器人的位置和姿态)
        # 从URDF文件: origin xyz="0.5 -0.04 0.57" rpy="0 0.314 0"
        self.camera_translation = np.array([0.5, -0.04, 0.57])  # 相机位置
        self.camera_pitch = 0.314  # 相机俯仰角 (弧度)

        rospy.Subscriber('/frontiers_2d', PointCloud2, self.frontier_callback)
        rospy.Subscriber('/magv/odometry/gt', Odometry, self.odom_callback)
        rospy.Subscriber('/magv/camera/image_compressed/compressed', CompressedImage, self.image_callback)
        rospy.Subscriber('/instruction', String, self.instruction_callback)
        self.waypoint_pub = rospy.Publisher('/way_point', PointStamped, queue_size=1)
        self.stop_pub = rospy.Publisher('/status', Int32, queue_size=1)
        # 添加定时器，按频率发布目标点
        self.waypoint_timer = rospy.Timer(rospy.Duration(1.0/self.waypoint_publish_rate), self.timer_callback)

    def timer_callback(self, event):
            
        current_time = rospy.Time.now()
        if (current_time - self.last_waypoint_time).to_sec() >= 1.0/self.waypoint_publish_rate:
            self.publish_current_waypoint()
            self.last_waypoint_time = current_time

    def frontier_callback(self, msg):
        if not self.exploration_active:
            return
        # 解析PointCloud2为frontier点列表
        all_frontiers = []
        for p in pc2.read_points(msg, field_names=("x", "y", "z"), skip_nans=True):
            all_frontiers.append(np.array([p[0], p[1], p[2]]))
        
        # 过滤出在探索区域内的frontiers
        self.frontiers = self.filter_frontiers_in_exploration_area(all_frontiers)
        
        # 检查探索区域内是否有frontiers
        if not self.frontiers:
            self.consecutive_empty_frontiers += 1
            rospy.logwarn(f"探索区域内没有发现frontiers，连续次数: {self.consecutive_empty_frontiers}")
            rospy.loginfo(f"探索区域: X[{self.exploration_bounds['x_min']:.1f}, {self.exploration_bounds['x_max']:.1f}], "
                         f"Y[{self.exploration_bounds['y_min']:.1f}, {self.exploration_bounds['y_max']:.1f}]")
            
            # 检查是否达到终止条件
            if self.consecutive_empty_frontiers >= self.max_empty_count:
                if self.exploration_active:
                    self.finish_exploration()
        else:
            # 重置计数器
            self.consecutive_empty_frontiers = 0
            #rospy.loginfo(f"探索区域内发现 {len(self.frontiers)} 个frontiers")
            #rospy.loginfo(f"frontiers: {self.frontiers}")
            
        # 如果没有目标点且探索仍在进行，选择一个新的目标点
        if self.current_goal is None and self.exploration_active and self.frontiers:
            rospy.loginfo("没有目标点，选择一个新的目标点")
            self.select_next_waypoint()

    def odom_callback(self, msg):
        self.robot_pos = np.array([
            msg.pose.pose.position.x,
            msg.pose.pose.position.y,
            msg.pose.pose.position.z
        ])
        
        # 提取机器人朝向（yaw角）
        orientation = msg.pose.pose.orientation
        self.robot_yaw = self.quaternion_to_yaw(orientation)
        
        # 设置起始位置（如果还没有设置）
        if self.robot_start_pos is None:
            self.robot_start_pos = self.robot_pos.copy()
            rospy.loginfo(f"设置机器人起始位置: {self.robot_start_pos}")
            rospy.loginfo(f"探索区域边界: X[{self.exploration_bounds['x_min']:.1f}, {self.exploration_bounds['x_max']:.1f}], "
                         f"Y[{self.exploration_bounds['y_min']:.1f}, {self.exploration_bounds['y_max']:.1f}]")
        
        # 检查是否到达目标点
        if self.current_goal is not None and self.robot_pos is not None:
            dist = np.linalg.norm(self.current_goal[:2] - self.robot_pos[:2])
            if dist < self.goal_reached_threshold:
                self.current_goal = None
                # 只有在探索仍在进行时才选择下一个目标点
                if self.exploration_active and self.frontiers:
                    rospy.loginfo("到达目标点，选择下一个目标点")
                    self.select_next_waypoint()
                else:
                    rospy.loginfo("探索已完成，不再选择新的目标点")
                    self.stop_pub.publish(0)

    def image_callback(self, msg):
        """处理摄像头图像，检测Aruco标记"""
        if self.aruco_detected or not self.exploration_active:
            return
            
        try:
            # 转换ROS图像到OpenCV格式
            cv_image = self.bridge.compressed_imgmsg_to_cv2(msg, "bgr8")
            #cv2.imshow("image", cv_image)
            # 检测Aruco标记
            corners, ids, _ = aruco.detectMarkers(cv_image, self.aruco_dict, parameters=self.aruco_params)
            if ids is not None and len(ids) > 0:
                rospy.loginfo(f"检测到Aruco标记，ID: {ids.flatten()}")
                
                # 估计姿态
                rvecs, tvecs, _ = aruco.estimatePoseSingleMarkers(
                    corners, self.aruco_size, self.camera_matrix, self.dist_coeffs
                )
                
                # 计算第一个检测到的Aruco在机器人坐标系中的位置
                if len(tvecs) > 0:
                    self.calculate_aruco_position(tvecs[0], rvecs[0])
                    
                    # 停止探索
                    self.stop_exploration()
                    
        except Exception as e:
            rospy.logwarn(f"图像处理出错: {e}")

    def calculate_aruco_position(self, tvec, rvec):
        """计算Aruco标记在机器人坐标系中的位置"""
        if self.robot_pos is None:
            return
            
        # tvec是相对于摄像头的位置 (x, y, z)
        # 需要转换到机器人坐标系
        
        # 摄像头坐标系: x-右, y-下, z-前
        # 机器人坐标系: x-前, y-左, z-上
        
        camera_x, camera_y, camera_z = tvec.flatten()
        
        # 创建相机坐标系中的位置向量
        camera_pos = np.array([camera_x, camera_y, camera_z])
        
        # 考虑相机的俯仰角进行坐标变换
        cos_pitch = math.cos(self.camera_pitch)
        sin_pitch = math.sin(self.camera_pitch)
        
        # 旋转矩阵 (绕Y轴旋转俯仰角)
        # 同时进行坐标系转换：摄像头坐标系到机器人坐标系
        # 摄像头z轴(前) -> 机器人x轴(前)
        # 摄像头x轴(右) -> 机器人y轴(左)，取负号
        # 摄像头y轴(下) -> 机器人z轴(上)，取负号
        
        robot_relative_x = camera_z * cos_pitch + camera_y * sin_pitch
        robot_relative_y = -camera_x  # 摄像头的x轴对应机器人的y轴(左)，取负号
        robot_relative_z = -camera_z * sin_pitch + camera_y * cos_pitch
        
        # 加上相机相对于机器人的位置偏移
        robot_relative_x += self.camera_translation[0]
        robot_relative_y += self.camera_translation[1]
        robot_relative_z += self.camera_translation[2]
        
        # 考虑机器人的朝向，转换到世界坐标系
        cos_yaw = math.cos(self.robot_yaw)
        sin_yaw = math.sin(self.robot_yaw)
        
        world_x = self.robot_pos[0] + robot_relative_x * cos_yaw - robot_relative_y * sin_yaw
        world_y = self.robot_pos[1] + robot_relative_x * sin_yaw + robot_relative_y * cos_yaw
        world_z = self.robot_pos[2] + robot_relative_z
        
        self.aruco_position = np.array([world_x, world_y, world_z])
        self.aruco_detected = True
        
        rospy.loginfo(f"Aruco标记位置 (世界坐标): {self.aruco_position}")
        rospy.loginfo(f"相对机器人距离: {math.sqrt(robot_relative_x**2 + robot_relative_y**2):.2f}m")

    def stop_exploration(self):
        """停止探索"""
        self.exploration_active = False
        self.current_goal = self.aruco_position
        rospy.loginfo("=" * 50)
        rospy.loginfo("检测到Aruco标记，停止探索！")
        if self.aruco_position is not None:
            rospy.loginfo(f"Aruco标记位置: {self.aruco_position}")
        rospy.loginfo("=" * 50)

    def select_next_waypoint(self):
        """选择下一个目标点"""
        if not self.exploration_active:
            rospy.loginfo("探索已完成，不再选择新的目标点")
            return
            
        if not self.frontiers or self.robot_pos is None:
            rospy.loginfo("探索区域内没有可用的frontier或机器人位置未知")
            return
            
        # 计算每个frontier的综合得分
        frontier_scores = []
        for frontier in self.frontiers:
            # 距离得分（距离越近得分越高）
            distance = np.linalg.norm(frontier - self.robot_pos)
            distance_score = 1.0 / (1.0 + distance)  # 避免除零
            
            # 方向得分
            direction_score = self.calculate_direction_score(frontier)
            
            # 综合得分 = 距离得分 + 方向权重 * 方向得分
            total_score = distance_score + self.direction_weight * direction_score
            frontier_scores.append(total_score)
        # 选择得分最高的frontier
        best_idx = np.argmax(frontier_scores)
        target = self.frontiers[best_idx]
        self.current_goal = target
        
        # 计算并显示选择信息
        distance = np.linalg.norm(target - self.robot_pos)
        direction_score = self.calculate_direction_score(target)
        rospy.loginfo(f"选择目标点: {target}, 距离: {distance:.2f}m, 方向得分: {direction_score:.2f}, 偏好方向: {self.preferred_direction}")

    def publish_current_waypoint(self):
            
        if self.current_goal is None:
            return

        # 发布目标点
        point_msg = PointStamped()
        point_msg.header.stamp = rospy.Time.now()
        point_msg.header.frame_id = "map"  # 根据你的TF设置，可能需要改为odom等
        point_msg.point.x = self.current_goal[0]
        point_msg.point.y = self.current_goal[1]
        point_msg.point.z = self.current_goal[2]
        self.waypoint_pub.publish(point_msg)
        
        # 显示发布信息
        if self.robot_pos is not None:
            distance = np.linalg.norm(self.current_goal - self.robot_pos)
            #rospy.loginfo(f"发布目标点: {self.current_goal}, 距离: {distance:.2f}m")

    def filter_frontiers_in_exploration_area(self, all_frontiers):
        """过滤出在探索区域内的frontiers"""
        if self.robot_start_pos is None:
            # 如果还没有设置起始位置，使用当前机器人位置作为参考
            if self.robot_pos is not None:
                self.robot_start_pos = self.robot_pos.copy()
            else:
                return all_frontiers  # 如果连机器人位置都不知道，返回所有frontiers
        
        filtered_frontiers = []
        for frontier in all_frontiers:
            # 计算相对于起始位置的坐标
            relative_pos = frontier - self.robot_start_pos
            # 检查是否在探索区域内
            if (self.exploration_bounds['x_min'] <= relative_pos[0] <= self.exploration_bounds['x_max'] and
                self.exploration_bounds['y_min'] <= relative_pos[1] <= self.exploration_bounds['y_max']):
                filtered_frontiers.append(frontier)
        
        return filtered_frontiers

    def finish_exploration(self):
        """结束探索"""
        self.exploration_active = False
        self.current_goal = None
        rospy.loginfo("=" * 50)
        rospy.loginfo("探索完成！指定探索区域内没有更多frontiers")
        rospy.loginfo(f"探索区域: X[{self.exploration_bounds['x_min']:.1f}, {self.exploration_bounds['x_max']:.1f}], "
                     f"Y[{self.exploration_bounds['y_min']:.1f}, {self.exploration_bounds['y_max']:.1f}]")
        rospy.loginfo("=" * 50)
        
        # 可以在这里添加其他结束探索的逻辑
        # 比如发布探索完成的消息，或者执行其他操作

    def calculate_direction_score(self, frontier_point):
        """计算frontier点相对于偏好方向的得分"""
        if self.robot_pos is None:
            return 0.0
            
        # 计算frontier相对于机器人的方向向量
        direction_vector = frontier_point - self.robot_pos
        direction_vector[2] = 0  # 忽略z轴，只考虑水平方向
        
        # 计算角度（相对于机器人前方）
        angle = math.degrees(math.atan2(direction_vector[1], direction_vector[0]))
        if angle < 0:
            angle += 360
            
        # 计算与偏好方向的角度差
        preferred_angle = self.direction_angles.get(self.preferred_direction, 0)
        angle_diff = abs(angle - preferred_angle)
        if angle_diff > 180:
            angle_diff = 360 - angle_diff
            
        # 计算方向得分（角度差越小得分越高）
        direction_score = 1.0 - (angle_diff / 180.0)
        return direction_score

    def set_preferred_direction(self, direction):
        """设置偏好方向"""
        if direction in self.direction_angles:
            self.preferred_direction = direction
            rospy.loginfo(f"设置偏好方向为: {direction}")
        else:
            rospy.logwarn(f"无效的方向: {direction}")

    def set_waypoint_publish_rate(self, rate):
        """设置目标点发布频率"""
        self.waypoint_publish_rate = rate
        # 重新设置定时器
        self.waypoint_timer.shutdown()
        self.waypoint_timer = rospy.Timer(rospy.Duration(1.0/self.waypoint_publish_rate), self.timer_callback)
        rospy.loginfo(f"设置目标点发布频率为: {rate} Hz")

    def quaternion_to_yaw(self, orientation):
        """将四元数转换为yaw角"""
        x = orientation.x
        y = orientation.y
        z = orientation.z
        w = orientation.w
        
        # 计算yaw角
        yaw = math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z))
        return yaw

    def instruction_callback(self, msg):
        rospy.loginfo(f"收到指令: {msg.data}")
        for i in range(10):
            direction = parse_goal_direction(msg.data)
            if validate_direction(direction):
                self.set_preferred_direction(direction)
                self.exploration_active = True
                break
            else:
                rospy.logwarn(f"无效的方向: {direction}, 尝试第{i+1}次")
        if not validate_direction(direction):
            rospy.logwarn("10次尝试失败，使用默认方向")
            self.set_preferred_direction(self.preferred_direction)
            self.exploration_active = True
        

if __name__ == '__main__':
    try:
        FrontierExplorer()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass

